{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1327, 20, 35, 25)\n",
      "(1327, 20, 35)\n",
      "(1327, 20, 35)\n",
      "Data Size  10201\n",
      "word_emb {'shape': (10201, 100)}\n",
      "_word_rnn {'shape': (100, 200)}\n",
      "word_hs {'shape': (200, 10201)}\n",
      "H_softmax, two level  101 101\n",
      "TEST COMPILE\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "os.environ['THEANO_FLAGS']='mode=FAST_RUN,device=gpu0,floatX=float32'\n",
    "import time\n",
    "import sys, getopt\n",
    "import numpy as NP\n",
    "import theano\n",
    "import theano.tensor as T\n",
    "import theano.tensor.nnet as NN\n",
    "from wrapper import *\n",
    "from collections import OrderedDict\n",
    "from txt_data_word import *\n",
    "\n",
    "gru2_dim = 200 # gru in word LM\n",
    "emb_dim = 100 # word emb dim, must equal to gru1_dim due to the high way network\n",
    "batch_size = 1\n",
    "drop_flag = False\n",
    "train_flag = True\n",
    "mode = None\n",
    "fname = 'ptb_vis'\n",
    "Dataset = txt_data('ptb.train', 'ptb.valid', 'ptb_dict', ['ptb.train', 'ptb.valid', 'ptb.test'], batch_size=20)\n",
    "\n",
    "word_dim = Dataset.w_dim # vocabulary size\n",
    "\n",
    "var_lr = theano.shared(NP.asarray(1.0, dtype=theano.config.floatX)) #learning rate of rmsprop\n",
    "var_rescale = theano.shared(NP.asarray(5.0, dtype=theano.config.floatX)) # rescale the gradient norm\n",
    "word_in = T.imatrix()\n",
    "word_target = T.imatrix()\n",
    "model = Model()\n",
    "\n",
    "#get the mask for handling various length in the batch, the all zeros vector would be masked out\n",
    "def get_mask(x_in, mask_value=0.):\n",
    "    return T.shape_padright(T.any((1. - T.eq(x_in, mask_value)), axis=-1))\n",
    "\n",
    "def _word_rnn(x, prev_h):\n",
    "    gru1 = model.gru(x, prev_h, name='_word_rnn', shape=(emb_dim, gru2_dim))\n",
    "    return gru1\n",
    "\n",
    "\n",
    "drop_flag = False\n",
    "EPSI = 1e-15\n",
    "\n",
    "train_state_h1 = theano.shared(NP.zeros((batch_size , gru2_dim), dtype=theano.config.floatX))\n",
    "\n",
    "test_state_h1 = theano.shared(NP.zeros((batch_size , gru2_dim), dtype=theano.config.floatX))\n",
    "\n",
    "def get_express(train=False, emb_flag=None):\n",
    "        global drop_flag\n",
    "        drop_flag = train\n",
    "        state_updates = []\n",
    "        state_h1 = train_state_h1 if train else test_state_h1\n",
    "        word_embs = model.embedding(word_in, name = 'word_emb', shape = (word_dim, emb_dim))\n",
    "        \n",
    "        seq_len = word_in.shape[1]\n",
    "        sc, _ = theano.scan(_word_rnn, sequences=[word_embs.dimshuffle(1,0,2)], outputs_info=[state_h1])\n",
    "        word_out = sc.dimshuffle(1,0,2)\n",
    "        state_updates.append((state_h1, sc[-1]))\n",
    "        if drop_flag:\n",
    "            word_out = drop5.drop(word_out)\n",
    "\n",
    "        word_out = word_out.reshape((word_out.shape[0]*(word_out.shape[1]), gru2_dim))\n",
    "\n",
    "        word_out = model.h_softmax(x_in = word_out, name='word_hs', shape=(gru2_dim, word_dim))\n",
    "        \n",
    "        return word_out[-1]\n",
    "\n",
    "wo = get_express(train=False, emb_flag=mode)\n",
    "predict_func = theano.function([word_in], [wo], allow_input_downcast=True)\n",
    "\n",
    "model.load('ptb_vis_19')\n",
    "print 'TEST COMPILE'\n",
    "    \n",
    "\n",
    "def reset_states():\n",
    "    test_state_h1.set_value(test_state_h1.get_value()*0.0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style>.container { width:100% !important; }</style>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['it', u'subject', u'will', u'normally', u'a', u'party', u'of', u'seven', '<unk>', u'payments', '<eos>']\n"
     ]
    }
   ],
   "source": [
    "from IPython.core.display import display, HTML\n",
    "display(HTML(\"<style>.container { width:100% !important; }</style>\"))\n",
    "import numpy.random as RNG\n",
    "Wdict = Dataset.Wdict\n",
    "num2w = Dataset.num2w\n",
    "num2w[0]='<unk>'\n",
    "st=['it'] # given the start words\n",
    "\n",
    "reset_states()\n",
    "\n",
    "for i in xrange(10): # the sentence length\n",
    "    w_in = []\n",
    "    for j in st:\n",
    "        w_in.append(Wdict[j])\n",
    "    t = predict_func(NP.asarray([w_in]))[0] # generating the next word\n",
    "    next_idx=RNG.choice(word_dim, 1, p=t)[0] # sampling\n",
    "    next_word = num2w[next_idx]\n",
    "    st.append(next_word)\n",
    "\n",
    "for i in xrange(len(st)):\n",
    "    if st[i] ==u'\\u1d94':\n",
    "        st[i] = '<unk>'\n",
    "    if st[i] ==u'\\u1d93':\n",
    "        st[i] = '<eos>'\n",
    "print st # generated sentence "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['THEANO_FLAGS']='mode=FAST_RUN,device=gpu0,floatX=float32'\n",
    "import time\n",
    "import sys, getopt\n",
    "import numpy as NP\n",
    "import theano\n",
    "import theano.tensor as T\n",
    "import theano.tensor.nnet as NN\n",
    "from wrapper import *\n",
    "from collections import OrderedDict\n",
    "from txt_data_word import *\n",
    "\n",
    "class Moving_AVG(object):\n",
    "\tdef __init__(self, array_size=500):\n",
    "\t\tself.array_size=array_size\n",
    "\t\tself.queue = NP.zeros((array_size, ))\n",
    "\t\tself.idx = 0\n",
    "\t\tself.filled = False\n",
    "\tdef append(self, x):\n",
    "\t\tself.queue[self.idx]=x\n",
    "\t\tself.idx = (self.idx + 1)%self.array_size\n",
    "\t\tif not self.filled and self.idx==0:\n",
    "\t\t\tself.filled = True\n",
    "\tdef get_avg(self):\n",
    "\t\tresult = NP.mean(self.queue) if self.filled else NP.sum(self.queue)/NP.sum(self.queue!=0)\n",
    "\t\treturn result\n",
    "\n",
    "\n",
    "\n",
    "gru2_dim = 200 # gru in word LM\n",
    "emb_dim = 100 # word emb dim, must equal to gru1_dim due to the high way network\n",
    "char_emb_dim = 15\n",
    "batch_size = 20\n",
    "drop_flag = False\n",
    "RHO=0.4\n",
    "half_W = 6\n",
    "FREQ_T = 1e-3\n",
    "train_flag = True\n",
    "mode = None\n",
    "fname = 'ptb_vis'\n",
    "Dataset = txt_data('ptb.train', 'ptb.valid', 'ptb_dict', ['ptb.train', 'ptb.valid', 'ptb.test'], batch_size=20)\n",
    "\n",
    "word_dim = Dataset.w_dim # word dict size\n",
    "\n",
    "var_lr = theano.shared(NP.asarray(1.0, dtype=theano.config.floatX)) #learning rate of rmsprop\n",
    "var_rescale = theano.shared(NP.asarray(5.0, dtype=theano.config.floatX)) # rescale the gradient norm\n",
    "word_in = T.imatrix()\n",
    "word_target = T.imatrix()\n",
    "model = Model()\n",
    "drop2 = Dropout(shape=(word_target.shape[0]*word_target.shape[1], emb_dim), prob=0.0)\n",
    "drop3 = Dropout(shape=(word_target.shape[0], gru2_dim), prob=0.0)\n",
    "drop4 = Dropout(shape=(word_target.shape[0], emb_dim), prob=0.0)\n",
    "drop5 = Dropout(shape=(word_target.shape[0], word_target.shape[1], gru2_dim), prob=0.5)\n",
    "drop6 = Dropout(shape=(word_target.shape[0], word_target.shape[1], gru2_dim), prob=0.0)\n",
    "\n",
    "def categorical_crossentropy(prob, true_idx):\n",
    "    true_idx = T.arange(true_idx.shape[0]) * word_dim + true_idx\n",
    "    t1 = prob.flatten()[true_idx]\n",
    "    return -t1\n",
    "\n",
    "#get mask for handle variable length in batch, all zeros vector will be masked out\n",
    "def get_mask(x_in, mask_value=0.):\n",
    "    return T.shape_padright(T.any((1. - T.eq(x_in, mask_value)), axis=-1))\n",
    "\n",
    "#character step\n",
    "\n",
    "def letanh(x):\n",
    "    return 1.7159*tanh(2.0/3.0*x)\n",
    "\n",
    "def relu(x):\n",
    "    return T.switch(x>0, x, 0.1*x)\n",
    "\n",
    "def high_way(x_in=None, name='', shape=[]):\n",
    "\tg = sigmoid(model.fc(x_in = x_in, name=name+'_g', shape=shape))\n",
    "\th = relu(model.fc(x_in=x_in, name=name+'_h', shape=shape))\n",
    "\treturn h*g + x_in * (1. - g)\n",
    "\n",
    "def softmax(x):\n",
    "    e_x = T.exp(x)\n",
    "    sm = e_x / e_x.sum(axis=1, keepdims=True)\n",
    "    return sm\n",
    "\n",
    "def softmax_t(x, t):\n",
    "    e_x = T.exp(1.0*x/t)\n",
    "    sm = e_x / e_x.sum(axis=1, keepdims=True)\n",
    "    return sm\n",
    "\n",
    "def log_softmax(x):\n",
    "    xdev = x - x.max(1, keepdims=True)\n",
    "    return xdev - T.log(T.sum(T.exp(xdev), axis=1, keepdims=True))\n",
    "\n",
    "def re_sigmoid(x):\n",
    "    sx = sigmoid(x)\n",
    "    return T.switch(sx>0.5, sx, 0.0)\n",
    "\n",
    "def _word_rnn(x, prev_h):\n",
    "    gru1 = model.gru(x, prev_h, name='_word_rnn', shape=(emb_dim, gru2_dim))\n",
    "    return gru1\n",
    "\n",
    "sc, _ = theano.scan(_word_rnn, sequences=[word_embs.dimshuffle(1,0,2)], /\n",
    "                            outputs_info=[state_h1])\n",
    "\n",
    "prev_h = 0\n",
    "for t in xrange(10):\n",
    "    gru1[t] = model.gru(x[t], prev_h, name='_word_rnn', shape=(emb_dim, gru2_dim))\n",
    "    prev_h = gru1[t]\n",
    "    \n",
    "drop_flag = False\n",
    "EPSI = 1e-15\n",
    "\n",
    "train_state_h1 = theano.shared(NP.zeros((batch_size , gru2_dim), dtype=theano.config.floatX))\n",
    "\n",
    "test_state_h1 = theano.shared(NP.zeros((batch_size , gru2_dim), dtype=theano.config.floatX))\n",
    "\n",
    "def get_express(train=False, emb_flag=None):\n",
    "        global drop_flag\n",
    "        drop_flag = train\n",
    "        state_updates = []\n",
    "        state_h1 = train_state_h1 if train else test_state_h1\n",
    "        word_embs = model.embedding(word_in, name = 'word_emb', shape = (word_dim, emb_dim))\n",
    "        \n",
    "        seq_len = word_in.shape[1]\n",
    "        sc, _ = theano.scan(_word_rnn, sequences=[word_embs.dimshuffle(1,0,2)], /\n",
    "                            outputs_info=[state_h1])\n",
    "        word_out = sc.dimshuffle(1,0,2)\n",
    "        state_updates.append((state_h1, sc[-1]))\n",
    "        if drop_flag:\n",
    "            word_out = drop5.drop(word_out)\n",
    "\n",
    "        word_out = word_out.reshape((word_target.shape[0]*(word_target.shape[1]), gru2_dim))\n",
    "        f_word_target = word_target.reshape((word_target.shape[0]*(word_target.shape[1]), ))\n",
    "\n",
    "        word_out = model.h_softmax(x_in = word_out, y_in = f_word_target, name='word_hs', /\n",
    "                                   shape=(gru2_dim, word_dim))\n",
    "        word_out = -T.log(T.clip(word_out, EPSI, 1.0-EPSI))\n",
    "\n",
    "        mask_PPL = (1. - T.eq(f_word_target, -1))*word_out\n",
    "        PPL_word_LM = T.sum(mask_PPL)\n",
    "        cost_word_LM = PPL_word_LM/T.sum(f_word_target>=0)\n",
    "        cost_all = cost_word_LM\n",
    "        if train:\n",
    "            grad_all, grad_norm = rmsprop(cost_all, model.get_params(), lr=var_lr, /\n",
    "                                          epsilon=var_lr**2, rescale = var_rescale , ignore_input_disconnect=True, return_norm=True)\n",
    "            return cost_all, PPL_word_LM, grad_all + state_updates+model.get_updates(), grad_norm\n",
    "        else:\n",
    "            return cost_all, PPL_word_LM, state_updates+model.get_updates()\n",
    "\n",
    "cost_all, PPL, all_updates = get_express(train=False, emb_flag=mode)\n",
    "test_func_hw = theano.function([word_in, word_target], [cost_all, PPL], updates=all_updates, /\n",
    "                               allow_input_downcast=True, on_unused_input='ignore')\n",
    "\n",
    "print 'TEST COMPILE'\n",
    "\n",
    "if train_flag :\n",
    "    cost_all, PPL, all_updates, grad_norm = get_express(train=True, emb_flag=mode)\n",
    "    train_func_hw = theano.function([word_in, word_target], [cost_all, PPL, grad_norm], / \n",
    "                                    updates=all_updates, allow_input_downcast=True, on_unused_input='ignore')\n",
    "    print 'TRAIN COMPILE'       \n",
    "\n",
    "train_func = None\n",
    "test_func = None\n",
    "\n",
    "def reset_states():\n",
    "    test_state_h1.set_value(test_state_h1.get_value()*0.0)\n",
    "\n",
    "last_valid_PPL = 2e10\n",
    "valid_increase_cnt = 0 #if loss on validation set didn't decrease in N steps, change the learning rate as half\n",
    "reset_states()\n",
    "for i in xrange(0, 20):\n",
    "    train_func = train_func_hw if train_flag else None\n",
    "    test_func = test_func_hw\n",
    "    if train_flag:\n",
    "        mytime = time.time()\n",
    "        train_batchs = int(Dataset.train_size/batch_size)\n",
    "        ma_cost = Moving_AVG(train_batchs) \n",
    "        for j in xrange(train_batchs):\n",
    "            word, word_label = Dataset.get_batch(only_word=True)\n",
    "            n_cost, n_ppl, n_norm = train_func(word, word_label)\n",
    "            ma_cost.append(n_cost)\n",
    "            if j%100==0:\n",
    "                print 'Epoch = ', str(i), ' Batch = ', str(j), ' Cost = ', n_cost, ' PPL = ', NP.exp(n_ppl/NP.sum(word_label>=0)), ' AVG Cost = ', ma_cost.get_avg(), 'LEN = ', NP.shape(word_label)[1], 'GRAD Norm = ', n_norm #NP.mean(NP.abs(model.wmatrix(name = 'word_emb_gain', shape=((1,))).get_value()))\n",
    "\n",
    "            if j%5000==4999:\n",
    "                model.save(fname+'_every')\n",
    "        newtime = time.time()\n",
    "        print 'One Epoch Time = ', newtime-mytime\n",
    "        mytime = newtime\n",
    "        model.save(fname+'_'+str(i))\n",
    "    if not train_flag:\n",
    "        model.load(fname+'_'+str(i))\n",
    "    reset_states()\n",
    "    model.clear_state()\n",
    "    Dataset.test_data_idx = 0\n",
    "    Dataset.test_len_idx = 0\n",
    "    test_wper = []\n",
    "    test_wcnt = []\n",
    "    test_batchs = 0\n",
    "    while test_batchs<Dataset.test_size:\n",
    "        word, word_label = Dataset.get_batch(test=True, only_word=True)\n",
    "        n_cost, n_ppl = test_func(word, word_label)\n",
    "        test_wper.append(n_ppl)\n",
    "        test_wcnt.append(NP.sum(word_label>=0))\n",
    "        test_batchs+=len(word_label)\n",
    "        print ' Test Progress = ', 1.0*test_batchs/Dataset.test_size, \n",
    "    valid_PPL = NP.exp(NP.sum(test_wper)/NP.sum(test_wcnt))\n",
    "    print '\\nEpoch = ', str(i), ' Test Word PPL = ', valid_PPL\n",
    "    if (valid_PPL>150 and valid_PPL+1 > last_valid_PPL) or (valid_PPL>last_valid_PPL):\n",
    "        valid_increase_cnt += 1\n",
    "        if valid_increase_cnt>=1:\n",
    "            print 'change learning rate', var_lr.get_value()*0.5\n",
    "            var_lr.set_value(var_lr.get_value()*0.5)\n",
    "            valid_increase_cnt = 0\n",
    "    last_valid_PPL = valid_PPL"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
